from flask import Flask, request, jsonify
from flask_cors import CORS
import pandas as pd
import numpy as np
import json
import os
from datetime import datetime
import logging

def safe_json_serialize(obj):
    """安全地将NumPy类型转换为JSON可序列化的类型"""
    if isinstance(obj, np.integer):
        return int(obj)
    elif isinstance(obj, np.floating):
        return float(obj)
    elif isinstance(obj, np.ndarray):
        return obj.tolist()
    elif isinstance(obj, pd.Timestamp):
        return obj.isoformat()
    elif hasattr(obj, 'item'):  # 处理其他NumPy标量类型
        return obj.item()
    else:
        return obj

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = Flask(__name__)
CORS(app)  # 允许跨域请求

# 配置上传文件夹
UPLOAD_FOLDER = 'uploads'
if not os.path.exists(UPLOAD_FOLDER):
    os.makedirs(UPLOAD_FOLDER)

app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB max file size

class CSVGraphProcessor:
    """CSV数据处理器，用于生成图谱节点数据"""
    
    def __init__(self):
        self.color_palette = [
            '#FF6B6B', '#4ECDC4', '#45B7D1', '#96CEB4', '#FFEAA7',
            '#DDA0DD', '#98D8C8', '#F7DC6F', '#BB8FCE', '#85C1E9',
            '#F8C471', '#82E0AA', '#F1948A', '#85C1E9', '#D7BDE2',
            '#A9DFBF', '#F9E79F', '#D5DBDB', '#AED6F1', '#A3E4D7'
        ]
    
    def process_csv(self, file_path):
        """处理CSV文件并生成图谱数据"""
        try:
            # 读取CSV文件
            df = pd.read_csv(file_path, encoding='utf-8')
            
            # 如果编码失败，尝试其他编码
            if df.empty:
                df = pd.read_csv(file_path, encoding='gbk')
            
            logger.info(f"成功读取CSV文件，共 {len(df)} 行，{len(df.columns)} 列")
            
            # 生成节点数据
            nodes = self._generate_nodes(df)
            
            # 生成连线数据（基于列之间的相关性）
            links = self._generate_links(df)
            
            # 生成统计信息
            stats = self._generate_stats(df)
            
            return {
                'nodes': nodes,
                'links': links,
                'stats': stats,
                'success': True,
                'message': 'CSV数据解析成功'
            }
            
        except Exception as e:
            logger.error(f"处理CSV文件时出错: {str(e)}")
            return {
                'nodes': [],
                'links': [],
                'stats': {},
                'success': False,
                'message': f'处理CSV文件时出错: {str(e)}'
            }
    
    def _generate_nodes(self, df):
        """生成节点数据"""
        nodes = []
        
        # 为每个列创建节点
        for i, column in enumerate(df.columns):
            # 计算列的基本统计信息
            column_data = df[column].dropna()
            
            node = {
                'id': column,
                'name': column,
                'group': self._get_column_group(column, column_data),
                'color': self.color_palette[i % len(self.color_palette)],
                'size': self._calculate_node_size(column_data),
                'stats': {
                    'total_count': int(len(column_data)),
                    'null_count': int(df[column].isnull().sum()),
                    'unique_count': int(column_data.nunique()),
                    'data_type': str(column_data.dtype)
                }
            }
            
            # 如果是数值型数据，添加更多统计信息
            if pd.api.types.is_numeric_dtype(column_data):
                node['stats'].update({
                    'mean': float(column_data.mean()) if not column_data.empty else 0,
                    'std': float(column_data.std()) if not column_data.empty else 0,
                    'min': float(column_data.min()) if not column_data.empty else 0,
                    'max': float(column_data.max()) if not column_data.empty else 0
                })
            
            nodes.append(node)
        
        return nodes
    
    def _generate_links(self, df):
        """生成连线数据（基于列之间的相关性和主列关系）"""
        links = []
        
        # 1. 生成主列与其他列的连线
        primary_links = self._generate_primary_column_links(df)
        links.extend(primary_links)
        
        # 2. 生成数值型列之间的相关性连线
        correlation_links = self._generate_correlation_links(df)
        links.extend(correlation_links)
        
        return links
    
    def _generate_primary_column_links(self, df):
        """生成主列与其他列的连线关系"""
        links = []
        
        if len(df.columns) < 2:
            logger.info("列数少于2列，跳过主列连线生成")
            return links
        
        # 确定主列（第一列或第二列）
        primary_column = self._determine_primary_column(df)
        
        if not primary_column:
            logger.info("未找到合适的主列，跳过主列连线生成")
            return links
        
        logger.info(f"确定主列: {primary_column}")
        
        # 为主列与其他列创建连线
        for column in df.columns:
            if column != primary_column:
                # 计算关系强度（基于数据相似性或相关性）
                strength = self._calculate_column_relationship(df, primary_column, column)
                
                link = {
                    'source': primary_column,
                    'target': column,
                    'value': strength,
                    'strength': strength,
                    'type': 'primary_relationship',
                    'color': '#000000'  # 黑色表示主列关系
                }
                links.append(link)
        
        logger.info(f"生成了 {len(links)} 条主列连线")
        return links
    
    def _determine_primary_column(self, df):
        """确定主列：如果第一列是序号列，则使用第二列，否则使用第一列"""
        if len(df.columns) == 0:
            return None
        
        first_column = df.columns[0]
        first_column_data = df[first_column].dropna()
        
        logger.info(f"检查第一列 '{first_column}' 是否为序号列")
        
        # 检查第一列是否为序号列
        if self._is_sequence_column(first_column, first_column_data):
            # 如果第一列是序号，使用第二列作为主列
            if len(df.columns) > 1:
                logger.info(f"第一列 '{first_column}' 是序号列，使用第二列 '{df.columns[1]}' 作为主列")
                return df.columns[1]
            else:
                logger.info("第一列是序号列，但没有其他列可用")
                return None
        else:
            # 否则使用第一列作为主列
            logger.info(f"第一列 '{first_column}' 不是序号列，使用它作为主列")
            return first_column
    
    def _is_sequence_column(self, column_name, column_data):
        """判断列是否为序号列"""
        column_lower = column_name.lower()
        
        # 基于列名判断
        sequence_keywords = ['id', 'index', '序号', '编号', 'no', 'number', 'seq', 'sequence']
        if any(keyword in column_lower for keyword in sequence_keywords):
            return True
        
        # 基于数据特征判断
        if len(column_data) == 0:
            return False
        
        # 检查是否为连续递增的整数
        try:
            # 尝试转换为数值
            numeric_data = pd.to_numeric(column_data, errors='coerce')
            if numeric_data.isna().sum() == 0:  # 全部可以转换为数值
                # 检查是否为连续递增
                sorted_data = numeric_data.sort_values()
                if len(sorted_data) > 1:
                    diff = sorted_data.diff().dropna()
                    # 如果差值基本为1，且没有重复值，则认为是序号列
                    if (diff == 1).sum() >= len(diff) * 0.8 and sorted_data.nunique() == len(sorted_data):
                        return True
        except:
            pass
        
        # 检查是否为简单的序号模式（1,2,3...或A,B,C...）
        if len(column_data) <= 10:  # 只对小数据集检查
            unique_values = column_data.unique()
            if len(unique_values) == len(column_data):  # 无重复值
                # 检查是否为简单的序号
                if all(str(val).isdigit() for val in unique_values):
                    return True
        
        return False
    
    def _calculate_column_relationship(self, df, col1, col2):
        """计算两个列之间的关系强度"""
        try:
            data1 = df[col1].dropna()
            data2 = df[col2].dropna()
            
            if len(data1) == 0 or len(data2) == 0:
                return 0.1  # 默认弱关系
            
            # 如果都是数值型，计算相关性
            if pd.api.types.is_numeric_dtype(data1) and pd.api.types.is_numeric_dtype(data2):
                correlation = data1.corr(data2)
                return float(abs(correlation)) if not pd.isna(correlation) else 0.1
            
            # 如果都是分类数据，计算卡方检验或互信息
            elif not pd.api.types.is_numeric_dtype(data1) and not pd.api.types.is_numeric_dtype(data2):
                # 简单的分类关系强度计算
                # 基于共同值的比例
                common_ratio = len(set(data1) & set(data2)) / max(len(set(data1)), len(set(data2)))
                return float(common_ratio * 0.5)  # 分类关系的强度相对较低
            
            else:
                # 混合类型，使用默认关系强度
                return 0.3
        
        except Exception as e:
            logger.warning(f"计算列关系时出错: {str(e)}")
            return 0.1
    
    def _generate_correlation_links(self, df):
        """生成数值型列之间的相关性连线"""
        links = []
        
        # 只对数值型列计算相关性
        numeric_columns = df.select_dtypes(include=[np.number]).columns
        
        if len(numeric_columns) < 2:
            return links
        
        # 计算相关性矩阵
        correlation_matrix = df[numeric_columns].corr()
        
        # 生成连线（只保留相关性较强的连接）
        threshold = 0.3  # 相关性阈值
        
        for i, col1 in enumerate(numeric_columns):
            for j, col2 in enumerate(numeric_columns):
                if i < j:  # 避免重复和自连接
                    correlation = correlation_matrix.loc[col1, col2]
                    
                    if abs(correlation) >= threshold:
                        link = {
                            'source': col1,
                            'target': col2,
                            'value': float(abs(correlation)),
                            'strength': float(abs(correlation)),
                            'type': 'correlation',
                            'color': '#999' if correlation > 0 else '#ff6b6b'
                        }
                        links.append(link)
        
        return links
    
    def _get_column_group(self, column_name, column_data):
        """根据列名和数据类型确定节点组"""
        column_lower = column_name.lower()
        
        # 基于列名关键词分组
        if any(keyword in column_lower for keyword in ['id', 'key', 'index']):
            return 'identifier'
        elif any(keyword in column_lower for keyword in ['name', 'title', 'label']):
            return 'name'
        elif any(keyword in column_lower for keyword in ['date', 'time', 'year', 'month', 'day']):
            return 'datetime'
        elif any(keyword in column_lower for keyword in ['price', 'cost', 'amount', 'money']):
            return 'financial'
        elif any(keyword in column_lower for keyword in ['count', 'number', 'quantity']):
            return 'count'
        elif pd.api.types.is_numeric_dtype(column_data):
            return 'numeric'
        else:
            return 'categorical'
    
    def _calculate_node_size(self, column_data):
        """计算节点大小（基于数据的重要性和数量）"""
        base_size = 20
        
        # 基于数据量调整大小
        data_ratio = len(column_data.dropna()) / len(column_data) if len(column_data) > 0 else 0
        size_multiplier = 1 + data_ratio * 0.5
        
        # 基于唯一值比例调整大小
        if len(column_data) > 0:
            uniqueness_ratio = column_data.nunique() / len(column_data)
            size_multiplier += uniqueness_ratio * 0.3
        
        return int(base_size * size_multiplier)
    
    def _generate_stats(self, df):
        """生成数据集的统计信息"""
        return {
            'total_rows': int(len(df)),
            'total_columns': int(len(df.columns)),
            'numeric_columns': int(len(df.select_dtypes(include=[np.number]).columns)),
            'categorical_columns': int(len(df.select_dtypes(include=['object']).columns)),
            'datetime_columns': int(len(df.select_dtypes(include=['datetime64']).columns)),
            'missing_data_percentage': float((df.isnull().sum().sum() / (len(df) * len(df.columns))) * 100),
            'processed_at': datetime.now().isoformat()
        }

# 创建处理器实例
processor = CSVGraphProcessor()

@app.route('/api/health', methods=['GET'])
def health_check():
    """健康检查端点"""
    return jsonify({
        'status': 'healthy',
        'message': 'Graph Backend API is running',
        'timestamp': datetime.now().isoformat()
    })

@app.route('/api/upload-csv', methods=['POST'])
def upload_csv():
    """上传并处理CSV文件"""
    try:
        if 'file' not in request.files:
            return jsonify({
                'success': False,
                'message': '没有上传文件'
            }), 400
        
        file = request.files['file']
        
        if file.filename == '':
            return jsonify({
                'success': False,
                'message': '没有选择文件'
            }), 400
        
        if not file.filename.lower().endswith('.csv'):
            return jsonify({
                'success': False,
                'message': '只支持CSV文件格式'
            }), 400
        
        # 保存文件
        filename = f"{datetime.now().strftime('%Y%m%d_%H%M%S')}_{file.filename}"
        file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(file_path)
        
        logger.info(f"文件已保存: {file_path}")
        
        # 处理CSV文件
        result = processor.process_csv(file_path)
        
        # 保留文件用于调试（不删除）
        logger.info(f"文件已保留用于调试: {file_path}")
        
        # 使用安全的JSON序列化
        safe_result = json.loads(json.dumps(result, default=safe_json_serialize))
        return jsonify(safe_result)
        
    except Exception as e:
        logger.error(f"上传文件时出错: {str(e)}")
        return jsonify({
            'success': False,
            'message': f'处理文件时出错: {str(e)}'
        }), 500

@app.route('/api/process-csv-data', methods=['POST'])
def process_csv_data():
    """直接处理CSV数据（不通过文件上传）"""
    try:
        data = request.get_json()
        
        if not data or 'csv_data' not in data:
            return jsonify({
                'success': False,
                'message': '请求数据格式错误，需要包含csv_data字段'
            }), 400
        
        # 将CSV字符串转换为DataFrame
        from io import StringIO
        csv_string = data['csv_data']
        df = pd.read_csv(StringIO(csv_string))
        
        logger.info(f"成功解析CSV数据，共 {len(df)} 行，{len(df.columns)} 列")
        
        # 生成节点数据
        nodes = processor._generate_nodes(df)
        links = processor._generate_links(df)
        stats = processor._generate_stats(df)
        
        result = {
            'nodes': nodes,
            'links': links,
            'stats': stats,
            'success': True,
            'message': 'CSV数据解析成功'
        }
        # 使用安全的JSON序列化
        safe_result = json.loads(json.dumps(result, default=safe_json_serialize))
        return jsonify(safe_result)
        
    except Exception as e:
        logger.error(f"处理CSV数据时出错: {str(e)}")
        return jsonify({
            'success': False,
            'message': f'处理CSV数据时出错: {str(e)}'
        }), 500

@app.route('/api/process-csv-with-relationships', methods=['POST'])
def process_csv_with_relationships():
    """处理CSV文件并应用用户定义的关系"""
    try:
        if 'file' not in request.files:
            return jsonify({
                'success': False,
                'message': '没有上传文件'
            }), 400
        
        file = request.files['file']
        
        if file.filename == '':
            return jsonify({
                'success': False,
                'message': '没有选择文件'
            }), 400
        
        if not file.filename.lower().endswith('.csv'):
            return jsonify({
                'success': False,
                'message': '只支持CSV文件格式'
            }), 400
        
        # 获取关系数据
        relationships = request.form.get('relationships', '[]')
        nodes = request.form.get('nodes', '[]')
        
        try:
            relationships_data = json.loads(relationships)
            nodes_data = json.loads(nodes)
        except json.JSONDecodeError as e:
            return jsonify({
                'success': False,
                'message': f'关系数据格式错误: {str(e)}'
            }), 400
        
        # 保存文件
        filename = f"{datetime.now().strftime('%Y%m%d_%H%M%S')}_{file.filename}"
        file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(file_path)
        
        logger.info(f"文件已保存: {file_path}")
        
        # 处理CSV文件
        result = processor.process_csv(file_path)
        
        # 应用用户定义的关系
        if result['success'] and relationships_data:
            # 使用用户定义的关系替换自动生成的关系
            result['links'] = relationships_data
            result['message'] = 'CSV数据解析成功，已应用用户定义的关系'
            logger.info(f"应用了 {len(relationships_data)} 个用户定义的关系")
        
        # 保留文件用于调试（不删除）
        logger.info(f"文件已保留用于调试: {file_path}")
        
        # 使用安全的JSON序列化
        safe_result = json.loads(json.dumps(result, default=safe_json_serialize))
        return jsonify(safe_result)
        
    except Exception as e:
        logger.error(f"处理CSV文件时出错: {str(e)}")
        return jsonify({
            'success': False,
            'message': f'处理文件时出错: {str(e)}'
        }), 500

@app.route('/api/sample-data', methods=['GET'])
def get_sample_data():
    """获取示例数据用于测试"""
    sample_data = {
        'nodes': [
            {
                'id': 'name',
                'name': '姓名',
                'group': 'name',
                'color': '#FF6B6B',
                'size': 25,
                'stats': {
                    'total_count': 100,
                    'null_count': 0,
                    'unique_count': 100,
                    'data_type': 'object'
                }
            },
            {
                'id': 'age',
                'name': '年龄',
                'group': 'numeric',
                'color': '#4ECDC4',
                'size': 30,
                'stats': {
                    'total_count': 100,
                    'null_count': 2,
                    'unique_count': 45,
                    'data_type': 'int64',
                    'mean': 35.5,
                    'std': 12.3,
                    'min': 18,
                    'max': 65
                }
            },
            {
                'id': 'salary',
                'name': '薪资',
                'group': 'financial',
                'color': '#45B7D1',
                'size': 35,
                'stats': {
                    'total_count': 100,
                    'null_count': 1,
                    'unique_count': 78,
                    'data_type': 'float64',
                    'mean': 8500.0,
                    'std': 2500.0,
                    'min': 3000.0,
                    'max': 15000.0
                }
            }
        ],
        'links': [
            {
                'source': 'age',
                'target': 'salary',
                'value': 0.65,
                'strength': 0.65,
                'type': 'correlation',
                'color': '#999'
            }
        ],
        'stats': {
            'total_rows': 100,
            'total_columns': 3,
            'numeric_columns': 2,
            'categorical_columns': 1,
            'datetime_columns': 0,
            'missing_data_percentage': 1.0,
            'processed_at': datetime.now().isoformat()
        },
        'success': True,
        'message': '示例数据生成成功'
    }
    
    return jsonify(sample_data)


if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=5000)
